#include "SDL3_image/SDL_image.h"
#include "Graphics/sdlgpu.h"
#include "Graphics/gpurenderer.h"


namespace Jin {

	GPUDevice::GPUDevice(SDLWindow& window, GPUShaderFormat format, bool debugMode, stringv deviceName) noexcept : window(window) {
		device = SDL_CreateGPUDevice(format, debugMode, deviceName.data());
        if (device) {
            if (!SDL_ClaimWindowForGPUDevice(device, window)) {
                SDL_DestroyGPUDevice(device);
                device = nullptr;
                return;
            }
            swapchainTextureFormat = SDL_GetGPUSwapchainTextureFormat(device, window);
		}
	}

	GPUDevice::~GPUDevice() noexcept {
        Release();
	}

    void GPUDevice::Release() noexcept {
        if (device) {
            SDL_ReleaseWindowFromGPUDevice(device, window);
            SDL_DestroyGPUDevice(device);
            device = nullptr;
        }
    }

    GPUSampler::GPUSampler(GPUDevice& device, const GPUSamplerCreateInfo* createInfo) noexcept : device(device) {
        GPUSamplerCreateInfo localCreateInfo{};
        const SDL_GPUSamplerCreateInfo& samplerInfo = createInfo ? *createInfo : localCreateInfo;
        if (!createInfo) {
            localCreateInfo.min_filter = SDL_GPU_FILTER_LINEAR;
            localCreateInfo.mag_filter = SDL_GPU_FILTER_LINEAR;
            localCreateInfo.mipmap_mode = SDL_GPU_SAMPLERMIPMAPMODE_LINEAR;
            localCreateInfo.address_mode_u = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
            localCreateInfo.address_mode_v = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
            localCreateInfo.address_mode_w = SDL_GPU_SAMPLERADDRESSMODE_CLAMP_TO_EDGE;
            //localCreateInfo.mip_lod_bias = 0.0f;
            //localCreateInfo.max_anisotropy = 1.0f;
            //localCreateInfo.compare_op = SDL_GPU_COMPAREOP_ALWAYS;
            //localCreateInfo.min_lod = -1000.0f;
            //localCreateInfo.max_lod = 1000.0f;
            //localCreateInfo.enable_anisotropy = false;
            //localCreateInfo.enable_compare = false;
        }
        sampler = SDL_CreateGPUSampler(device, &samplerInfo);
    }

    GPUTexture::GPUTexture(GPUDevice& device, const GPUTextureCreateInfo& createInfo) noexcept : device(device), info(createInfo) {
        texture = SDL_CreateGPUTexture(device, &createInfo);
        if (texture) {
            viewport.w = (float)info.width;
            viewport.h = (float)info.height;
            scissor.size = { (int)info.width, (int)info.height };
            colorModType = COLORMODTYPE_MUL;
        }
    }

    GPUTexture::GPUTexture(GPUDevice& device, Uint32 width, Uint32 height, GPUTextureUsageFlags usage, GPUTextureFormat textureFormat) noexcept : device(device) {
        if (textureFormat == SDL_GPU_TEXTUREFORMAT_INVALID) {
            textureFormat = device.swapchainTextureFormat;
        }
        info.type = SDL_GPU_TEXTURETYPE_2D;
        info.format = textureFormat;
        info.usage = usage;
        info.width = width;
        info.height = height;
        info.layer_count_or_depth = 1;
        info.num_levels = 1;
        info.sample_count = SDL_GPU_SAMPLECOUNT_1;
        texture = SDL_CreateGPUTexture(device, &info);
        if (texture) {
            viewport.w = (float)info.width;
            viewport.h = (float)info.height;
            scissor.size = { (int)info.width, (int)info.height };
            colorModType = COLORMODTYPE_MUL;
        }
    }

    GPUTexture::GPUTexture(GPUDevice& device, stringv imageFile, GPUTextureUsageFlags usage, GPUTextureFormat textureFormat) noexcept : device(device) {
        if (textureFormat == SDL_GPU_TEXTUREFORMAT_INVALID) {
            textureFormat = device.swapchainTextureFormat;
        }
        SDLSurface img{ IMG_Load(imageFile.data()) };
        if (!img || !img.ConvertToFormat(GPUTextureFormatToSDLPixelFormat(textureFormat))) {
            return;
        }

        GPUTransferBuffer transferBuffer{ device, SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD, (Uint32)(img.Height() * img.Pitch()) };
        if (!transferBuffer) {
            return;
        }
        memcpy(GPUTransferBufferMapper{ transferBuffer }, img.GetPixels(), (size_t)(img.Height() * img.Pitch()));

        info.type = SDL_GPU_TEXTURETYPE_2D;
        info.format = textureFormat;
        info.usage = usage;
        info.width = (Uint32)img.Width();
        info.height = (Uint32)img.Height();
        info.layer_count_or_depth = 1;
        info.num_levels = 1;
        info.sample_count = SDL_GPU_SAMPLECOUNT_1;
        texture = SDL_CreateGPUTexture(device, &info);
        if (!texture) {
            return;
        }
        viewport.w = (float)info.width;
        viewport.h = (float)info.height;
        scissor.size = { (int)info.width, (int)info.height };
        colorModType = COLORMODTYPE_MUL;

        GPUCommandBuffer cmd{ device };
        if (!cmd) {
            return;
        }
        GPUTextureTransferInfo srcTexInfo{};
        srcTexInfo.transfer_buffer = transferBuffer;
        srcTexInfo.offset = 0;
        srcTexInfo.pixels_per_row = (Uint32)(img.Width());
        srcTexInfo.rows_per_layer = (Uint32)(img.Height());
        GPUTextureRegion dstTexRegion{};
        dstTexRegion.texture = texture;
        dstTexRegion.mip_level = 0;
        dstTexRegion.layer = 0;
        dstTexRegion.x = 0;
        dstTexRegion.y = 0;
        dstTexRegion.z = 0;
        dstTexRegion.w = (Uint32)(img.Width());
        dstTexRegion.h = (Uint32)(img.Height());
        dstTexRegion.d = 1;
        GPUCopyPass{ cmd }.UploadToGPUTexture(srcTexInfo, dstTexRegion);
    }

    GPUShader::GPUShader(GPUDevice& device, const GPUShaderInfo& shaderInfo) noexcept : device(device) {
        GPUShaderCreateInfo localCreateInfo{
        .code_size = shaderInfo.code.size(),
        .code = shaderInfo.code.data(),
        .entrypoint = shaderInfo.entrypoint.data(),
        .format = shaderInfo.format,
        .stage = shaderInfo.stage,
        .num_samplers = shaderInfo.numSamplers,
        .num_storage_textures = shaderInfo.numStorageTextures,
        .num_storage_buffers = shaderInfo.numStorageBuffers,
        .num_uniform_buffers = shaderInfo.numUniformBuffers
        };
        shader = SDL_CreateGPUShader(device, &localCreateInfo);
    }

    GPUGraphicsPipeline::GPUGraphicsPipeline(GPUDevice& device, GPUPrimitiveType type, const GPUVertexShader& vertexShader, const GPUFragmentShader& fragmentShader,
        const GPUColorTargetDescriptionsOpt& colorTargetDescriptions, const GPUGraphicsPipelineTargetInfoOpt& targetInfo,
        const GPURasterizerStateOpt& rasterizerState, const GPUMultisampleStateOpt& multisampleState, const GPUDepthStencilStateOpt& depthStencilState) noexcept : device(device) {
        
        GPUGraphicsPipelineCreateInfo localCreateInfo{};
        localCreateInfo.vertex_shader = vertexShader;
        localCreateInfo.fragment_shader = fragmentShader;
        localCreateInfo.vertex_input_state.vertex_buffer_descriptions = vertexShader.vertexBufferDescriptions.data();
        localCreateInfo.vertex_input_state.num_vertex_buffers = (Uint32)vertexShader.vertexBufferDescriptions.size();
        localCreateInfo.vertex_input_state.vertex_attributes = vertexShader.vertexAttributes.data();
        localCreateInfo.vertex_input_state.num_vertex_attributes = (Uint32)vertexShader.vertexAttributes.size();
        localCreateInfo.primitive_type = type;
        if (rasterizerState) {
            localCreateInfo.rasterizer_state = rasterizerState.value();
        }
        if (multisampleState) {
            localCreateInfo.multisample_state = multisampleState.value();
        }
        if (depthStencilState) {
            localCreateInfo.depth_stencil_state = depthStencilState.value();
        }
        if (targetInfo) {
            localCreateInfo.target_info = targetInfo.value();
        }
        if (colorTargetDescriptions) {
            localCreateInfo.target_info.color_target_descriptions = colorTargetDescriptions.value().data();
            localCreateInfo.target_info.num_color_targets = (Uint32)colorTargetDescriptions.value().size();
        }
        SDL_GPUColorTargetBlendState blendState{};
        blendState.enable_blend = true;
        blendState.src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA;
        blendState.dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
        blendState.color_blend_op = SDL_GPU_BLENDOP_ADD;
        blendState.src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE;
        blendState.dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA;
        blendState.alpha_blend_op = SDL_GPU_BLENDOP_ADD;
        blendState.color_write_mask = SDL_GPU_COLORCOMPONENT_R | SDL_GPU_COLORCOMPONENT_G | SDL_GPU_COLORCOMPONENT_B | SDL_GPU_COLORCOMPONENT_A;
        SDL_GPUColorTargetDescription defaultColorTargetDescriptions[1]{};
        defaultColorTargetDescriptions[0].format = device.swapchainTextureFormat;
        defaultColorTargetDescriptions[0].blend_state = blendState;
        if (!localCreateInfo.target_info.num_color_targets) {
            localCreateInfo.target_info.color_target_descriptions = defaultColorTargetDescriptions;
            localCreateInfo.target_info.num_color_targets = std::size(defaultColorTargetDescriptions);
        }
        pipeline = SDL_CreateGPUGraphicsPipeline(device, &localCreateInfo);
    }

    GPUComputePipeline::GPUComputePipeline(GPUDevice& device, const GPUComputeShaderInfo& shaderInfo) noexcept : device(device) {
        GPUComputePipelineCreateInfo localCreateInfo{
        .code_size = shaderInfo.code.size(),
        .code = shaderInfo.code.data(),
        .entrypoint = shaderInfo.entrypoint.data(),
        .format = shaderInfo.format,
        .num_samplers = shaderInfo.numSamplers,
        .num_readonly_storage_textures = shaderInfo.numReadonlyStorageTextures,
        .num_readonly_storage_buffers = shaderInfo.numReadonlyStorageBuffers,
        .num_readwrite_storage_textures = shaderInfo.numReadwriteStorageTextures,
        .num_readwrite_storage_buffers = shaderInfo.numReadwriteStorageBuffers,
        .num_uniform_buffers = shaderInfo.numUniformBuffers,
        .threadcount_x = shaderInfo.countX,
        .threadcount_y = shaderInfo.countY,
        .threadcount_z = shaderInfo.countZ
        };
        pipeline = SDL_CreateGPUComputePipeline(device, &localCreateInfo);
    }
}